{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/datadriven/anaconda3/envs/vllm/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "2024-09-28 14:57:39,057\tINFO util.py:154 -- Missing packages: ['ipywidgets']. Run `pip install -U ipywidgets`, then restart the notebook server for rich notebook output.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO 09-28 14:57:41 config.py:1559] Downcasting torch.float32 to torch.float16.\n",
      "WARNING 09-28 14:57:41 utils.py:721] Gemma 2 uses sliding window attention for every odd layer, which is currently not supported by vLLM. Disabling sliding window and capping the max length to the sliding window size (4096).\n",
      "INFO 09-28 14:57:41 llm_engine.py:184] Initializing an LLM engine (v0.5.5) with config: model='google/gemma-2-2B', speculative_config=None, tokenizer='google/gemma-2-2B', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, rope_scaling=None, rope_theta=None, tokenizer_revision=None, trust_remote_code=False, dtype=torch.float16, max_seq_len=4096, download_dir=None, load_format=LoadFormat.AUTO, tensor_parallel_size=1, pipeline_parallel_size=1, disable_custom_all_reduce=False, quantization=None, enforce_eager=False, kv_cache_dtype=auto, quantization_param_path=None, device_config=cuda, decoding_config=DecodingConfig(guided_decoding_backend='outlines'), observability_config=ObservabilityConfig(otlp_traces_endpoint=None, collect_model_forward_time=False, collect_model_execute_time=False), seed=0, served_model_name=google/gemma-2-2B, use_v2_block_manager=False, enable_prefix_caching=False)\n",
      "INFO 09-28 14:57:44 model_runner.py:879] Starting to load model google/gemma-2-2B...\n",
      "INFO 09-28 14:57:46 weight_utils.py:236] Using model weights format ['*.safetensors']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading safetensors checkpoint shards:   0% Completed | 0/3 [00:00<?, ?it/s]\n",
      "Loading safetensors checkpoint shards:  33% Completed | 1/3 [00:00<00:01,  1.81it/s]\n",
      "Loading safetensors checkpoint shards:  67% Completed | 2/3 [00:00<00:00,  3.03it/s]\n",
      "Loading safetensors checkpoint shards: 100% Completed | 3/3 [00:01<00:00,  2.50it/s]\n",
      "Loading safetensors checkpoint shards: 100% Completed | 3/3 [00:01<00:00,  2.48it/s]\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO 09-28 14:57:48 model_runner.py:890] Loading model weights took 4.8999 GB\n",
      "INFO 09-28 14:57:50 gpu_executor.py:121] # GPU blocks: 39897, # CPU blocks: 2520\n",
      "INFO 09-28 14:57:52 model_runner.py:1181] Capturing the model for CUDA graphs. This may lead to unexpected consequences if the model is not static. To run the model in eager mode, set 'enforce_eager=True' or use '--enforce-eager' in the CLI.\n",
      "INFO 09-28 14:57:52 model_runner.py:1185] CUDA graphs can take additional 1~3 GiB memory per GPU. If you are running out of memory, consider decreasing `gpu_memory_utilization` or enforcing eager mode. You can also reduce the `max_num_seqs` as needed to decrease memory usage.\n",
      "INFO 09-28 14:58:03 model_runner.py:1300] Graph capturing finished in 11 secs.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[RequestOutput(request_id=0, prompt='아이의 아버지인 외과의사가 말했어. \"난 수술 못해! 이 아이는 내 아들이라고!\" 이 외과의사는 소년에게 누구일까요?', prompt_token_ids=[2, 152384, 236137, 23745, 238364, 236183, 236589, 136201, 237233, 236137, 236417, 236361, 72163, 238267, 236770, 235265, 664, 240198, 22618, 239718, 145971, 237138, 235341, 11464, 72797, 236214, 58272, 23745, 94525, 112778, 4496, 11464, 136201, 237233, 236137, 236417, 236214, 44997, 237029, 70471, 130673, 237302, 236666, 239303, 237526, 235336], encoder_prompt=None, encoder_prompt_token_ids=None, prompt_logprobs=None, outputs=[CompletionOutput(index=0, text=' 사회 어떻게 그 아이를 심하게 괴롭혔을까요', token_ids=array('l', [208667, 53355, 242950, 237458, 20350, 72797, 236791, 135898, 101969, 235248, 242453, 243130, 245849, 236392, 239303, 237526]), cumulative_logprob=None, logprobs=None, finish_reason=length, stop_reason=None)], finished=True, metrics=RequestMetrics(arrival_time=1727535483.7901957, last_token_time=1727535483.7901957, first_scheduled_time=1727535483.8003397, first_token_time=1727535483.8159246, time_in_queue=0.01014399528503418, finished_time=1727535483.8969445, scheduler_time=0.0009138900013567763, model_forward_time=None, model_execute_time=None), lora_request=None)]"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "import re\n",
    "from typing import List, Dict, Optional\n",
    "from dataclasses import dataclass\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "from vllm import LLM, SamplingParams\n",
    "\n",
    "model = LLM(model=\"google/gemma-2-2B\", dtype='float16')\n",
    "prompt = '아이의 아버지인 외과의사가 말했어. \"난 수술 못해! 이 아이는 내 아들이라고!\" 이 외과의사는 소년에게 누구일까요?'\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "' 그대가 알지 못하던 외과의사가 한 것을 안'"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.generate(prompt, use_tqdm=False)[0].outputs[0].text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/datadriven/anaconda3/envs/vllm/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "2024-09-28 14:54:53,063\tINFO util.py:154 -- Missing packages: ['ipywidgets']. Run `pip install -U ipywidgets`, then restart the notebook server for rich notebook output.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO 09-28 14:54:55 config.py:1559] Downcasting torch.float32 to torch.float16.\n",
      "WARNING 09-28 14:54:55 utils.py:721] Gemma 2 uses sliding window attention for every odd layer, which is currently not supported by vLLM. Disabling sliding window and capping the max length to the sliding window size (4096).\n",
      "INFO 09-28 14:54:55 llm_engine.py:184] Initializing an LLM engine (v0.5.5) with config: model='google/gemma-2-2B', speculative_config=None, tokenizer='google/gemma-2-2B', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, rope_scaling=None, rope_theta=None, tokenizer_revision=None, trust_remote_code=False, dtype=torch.float16, max_seq_len=4096, download_dir=None, load_format=LoadFormat.AUTO, tensor_parallel_size=1, pipeline_parallel_size=1, disable_custom_all_reduce=False, quantization=None, enforce_eager=False, kv_cache_dtype=auto, quantization_param_path=None, device_config=cuda, decoding_config=DecodingConfig(guided_decoding_backend='outlines'), observability_config=ObservabilityConfig(otlp_traces_endpoint=None, collect_model_forward_time=False, collect_model_execute_time=False), seed=0, served_model_name=google/gemma-2-2B, use_v2_block_manager=False, enable_prefix_caching=False)\n",
      "INFO 09-28 14:54:58 model_runner.py:879] Starting to load model google/gemma-2-2B...\n",
      "INFO 09-28 14:55:00 weight_utils.py:236] Using model weights format ['*.safetensors']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading safetensors checkpoint shards:   0% Completed | 0/3 [00:00<?, ?it/s]\n",
      "Loading safetensors checkpoint shards:  33% Completed | 1/3 [00:00<00:01,  1.81it/s]\n",
      "Loading safetensors checkpoint shards:  67% Completed | 2/3 [00:00<00:00,  3.35it/s]\n",
      "Loading safetensors checkpoint shards: 100% Completed | 3/3 [00:01<00:00,  2.53it/s]\n",
      "Loading safetensors checkpoint shards: 100% Completed | 3/3 [00:01<00:00,  2.53it/s]\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO 09-28 14:55:02 model_runner.py:890] Loading model weights took 4.8999 GB\n",
      "INFO 09-28 14:55:04 gpu_executor.py:121] # GPU blocks: 39897, # CPU blocks: 2520\n",
      "INFO 09-28 14:55:06 model_runner.py:1181] Capturing the model for CUDA graphs. This may lead to unexpected consequences if the model is not static. To run the model in eager mode, set 'enforce_eager=True' or use '--enforce-eager' in the CLI.\n",
      "INFO 09-28 14:55:06 model_runner.py:1185] CUDA graphs can take additional 1~3 GiB memory per GPU. If you are running out of memory, consider decreasing `gpu_memory_utilization` or enforcing eager mode. You can also reduce the `max_num_seqs` as needed to decrease memory usage.\n",
      "INFO 09-28 14:55:17 model_runner.py:1300] Graph capturing finished in 11 secs.\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import re\n",
    "from typing import List, Dict, Optional\n",
    "from dataclasses import dataclass\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "from vllm import LLM, SamplingParams\n",
    "\n",
    "@dataclass\n",
    "class Path:\n",
    "    reasoning_text: str\n",
    "    score: float\n",
    "    answer_span: str\n",
    "    num_path: int\n",
    "\n",
    "@dataclass\n",
    "class DecodingInfo:\n",
    "    question: str\n",
    "    paths: List[Path]\n",
    "\n",
    "class CoTDecoder:\n",
    "    \"\"\"논문에서는 greedy decoding 대신 top-k를 이용하여 다양한 경로를 탐색하는 것을 권장합니다.\n",
    "    특히 각 경로에서 어떻게 생각하는지 평가할 수 있도록 여러 경로에서 다양한 토큰을 샘플링 해야한다고 설명합니다.\"\"\"\n",
    "    def __init__(self, model_name: str, \n",
    "                 device: str = 'cuda', \n",
    "                 max_new_tokens: int = 100, \n",
    "                 topk: int = 5, \n",
    "                 stop: List[str] = ['\\n\\n질문', '질문', 'Q:', '\\n\\nQ:', '\\n\\nExercise'],\n",
    "                 prompt: str = '', \n",
    "                 pattern: str = r'[가-힣a-zA-Z0-9\\s]+'):\n",
    "        self.model = LLM(model=model_name, dtype='float16')\n",
    "        self.tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "        self.device = device\n",
    "        self.max_new_tokens = max_new_tokens\n",
    "        self.stop = stop\n",
    "        self.topk = topk\n",
    "        self.model.llm_engine.model_config.max_logprobs = self.topk + 1\n",
    "        self.prompt = prompt\n",
    "        self.pattern = pattern\n",
    "\n",
    "    def search_cots(self, raw_prompt: str) -> DecodingInfo:\n",
    "        # 질문과 답변의 형태로 포맷을 변경합니다. \n",
    "        formatted_prompt = self.format_prompt(raw_prompt)\n",
    "        # 질문에 이어 생성된 단어를 top_k개만큼 생성하고, token_id, 생성된 토큰, 확률값을 저장한 topk_token을 생성합니다.\n",
    "        topk_tokens = self.get_first_topk_tokens(formatted_prompt)\n",
    "        # 생성된 토큰과 질문을 각각 합쳐주고 prompts라는 리스트에 저장합니다. \n",
    "        # 1개의 세트 예시  \n",
    "        # 질문:아이의 아버지인 외과의사가 말했어. \"난 수술 못해! 이 아이는 내 아들이라고!\" 이 외과의사는 소년에게 누구일까요?\n",
    "        # 답변:아이\n",
    "        prompts = [formatted_prompt + token for token in topk_tokens['decoded']]\n",
    "        outputs = self.generate_paths(prompts)\n",
    "        return self.calculate_score(raw_prompt, topk_tokens, outputs)\n",
    "\n",
    "    @torch.inference_mode()\n",
    "    def get_first_topk_tokens(self, prompt: str) -> Dict[str, List]:\n",
    "        sampling_params = SamplingParams(n=1, temperature=0, top_p=1, max_tokens=1, logprobs=self.topk, stop=self.stop)\n",
    "        # 모델이 입력된 prompt을 기준으로 prompt 뒤에 올 10개의 단어를 예측합니다. \n",
    "        outputs = self.model.generate(prompt, sampling_params, use_tqdm=False)[0].outputs[0].logprobs[0]\n",
    "        # decoded는 \"그\", \"\\n\\n\", \"소\", \"아이\", \"이\", \"아\", \"어\" 등의 단어가 생성되어 저장되어 있습니다. \n",
    "        # probs는 -2.064455270767212, -3.392580270767212 등의 로그 확률이 저정되어 있습니다. \n",
    "        # token_id는 당연히 token_id가 저장되어 있습니다. \n",
    "        topk_tokens = {'decoded': [], 'probs': [], 'token_id': [], 'logprobs': []}\n",
    "        for token_id, logprob_obj in outputs.items():\n",
    "            topk_tokens['logprobs'].append({token_id: logprob_obj})\n",
    "            topk_tokens['decoded'].append(logprob_obj.decoded_token)\n",
    "            topk_tokens['probs'].append(logprob_obj.logprob)\n",
    "            topk_tokens['token_id'].append(token_id)\n",
    "\n",
    "        # 로그 확률을 실제 확률로 변환합니다. \n",
    "        topk_tokens['probs'] = torch.exp(torch.tensor(topk_tokens['probs'])).tolist()\n",
    "        return topk_tokens\n",
    "\n",
    "    @torch.inference_mode()\n",
    "    def generate_paths(self, prompts: List[str]) -> Dict[int, Dict]:\n",
    "        # 논문에서는 top-k개의 토큰을 기반으로 다양한 경로를 생성해야한다고 설명합니다. 그 과정을 코드로 구현하였습니다. \n",
    "        # 리스트를 입력 받으면 배치로 생성이 됩니다. \n",
    "        sampling_params = SamplingParams(n=1, temperature=0, top_p=1, max_tokens=self.max_new_tokens, logprobs=2, stop=self.stop)\n",
    "        return self.model.generate(prompts, sampling_params, use_tqdm=False)\n",
    "\n",
    "    def format_prompt(self, raw_prompt: str) -> str:\n",
    "        return f'질문:{raw_prompt}\\n답변:{self.prompt}'\n",
    "\n",
    "\n",
    "    def calculate_score(self, prompt: str, topk_tokens: Dict, outputs: Dict) -> DecodingInfo:\n",
    "        paths = []\n",
    "        for k, output in enumerate(outputs):\n",
    "            reasoning = topk_tokens['decoded'][k] + output.outputs[0].text\n",
    "            reasoning = reasoning.strip()\n",
    "            \n",
    "            # 질문과 reasoning 간의 유사도를 계산 (간단한 방식으로 질문이 포함되었는지 확인)\n",
    "            question_similarity = self.calculate_question_similarity(prompt, reasoning)\n",
    "            \n",
    "            encode = self.tokenizer(reasoning, return_offsets_mapping=True)\n",
    "            answer_span = re.findall(self.pattern, reasoning)\n",
    "            \n",
    "            score = 0\n",
    "            if len(answer_span):\n",
    "                answer_span = answer_span[-1]\n",
    "                last_pattern_span = (reasoning.rfind(answer_span), reasoning.rfind(answer_span) + len(answer_span))\n",
    "                idx_answer = [i for i, span in enumerate(encode.offset_mapping)\n",
    "                            if (span[0] >= last_pattern_span[0] and span[1] <= last_pattern_span[1]) or\n",
    "                                (span[0] <= last_pattern_span[0] and span[1] >= last_pattern_span[1]) or\n",
    "                                (span[0] <= last_pattern_span[0] and span[1] > last_pattern_span[0])]\n",
    "\n",
    "                token_id = [encode.input_ids[idx] for idx in idx_answer]\n",
    "                output.outputs[0].logprobs.insert(0, topk_tokens['logprobs'][k])\n",
    "                filtered_answer = [output for i, output in enumerate(output.outputs[0].logprobs) if i in idx_answer]\n",
    "\n",
    "                sum_answer_span_probs = 0\n",
    "                for logprob_dict in filtered_answer:\n",
    "                    logprob_list = list(logprob_dict.items())\n",
    "                    if len(logprob_list) == 2:\n",
    "                        prob_diff = (torch.exp(torch.tensor([logprob_list[0][1].logprob])) - torch.exp(torch.tensor([logprob_list[1][1].logprob]))).item()\n",
    "                    else:\n",
    "                        prob_diff = torch.exp(torch.tensor([logprob_list[0][1].logprob])).item()\n",
    "                    sum_answer_span_probs += prob_diff\n",
    "                \n",
    "                # 질문과 비슷한 답변일 경우 페널티 적용\n",
    "                if question_similarity > 0.5:  # 질문과의 유사도가 높을수록 점수를 낮추기 위해 0.5 이상의 유사도에 패널티 적용\n",
    "                    sum_answer_span_probs *= (1 - question_similarity)  # 유사도가 높을수록 점수를 감소시킴\n",
    "\n",
    "                score = 0 if len(filtered_answer) == 0 else sum_answer_span_probs / len(filtered_answer)\n",
    "                answer_span = self.tokenizer.decode(token_id, skip_special_tokens=True).strip()\n",
    "            else:\n",
    "                answer_span = '|<NotFound>|'\n",
    "\n",
    "            paths.append(Path(reasoning_text=reasoning, \n",
    "                            score=score,\n",
    "                            answer_span=answer_span,\n",
    "                            num_path=k))\n",
    "\n",
    "        return DecodingInfo(question=prompt, paths=paths)\n",
    "\n",
    "    # 질문과 Reasoning의 유사도 계산하는 함수\n",
    "    def calculate_question_similarity(self, question: str, reasoning: str) -> float:\n",
    "        \"\"\" 질문과 reasoning 간의 유사도를 계산하는 간단한 함수. 유사도가 높으면 패널티를 부여한다 \"\"\"\n",
    "        question_words = set(question.split())\n",
    "        reasoning_words = set(reasoning.split())\n",
    "        \n",
    "        # 질문과 reasoning 간에 공통된 단어의 비율 계산\n",
    "        common_words = question_words.intersection(reasoning_words)\n",
    "        similarity = len(common_words) / len(question_words) if question_words else 0\n",
    "        \n",
    "        return similarity\n",
    "model_name = \"google/gemma-2-2B\"  \n",
    "# model_name = \"meta-llama/Llama-3.2-3B\"\n",
    "decoder = CoTDecoder(model_name)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Question: 아이의 아버지인 외과의사가 말했어. \"난 수술 못해! 이 아이는 내 아들이라고!\" 이 외과의사는 소년에게 누구일까요?\n",
      "Path 0:\n",
      "  Reasoning: 아버지인 외과의사가 말했어. \"난 수술 못해! 이 아이는 내 아들이라고!\" 이 외과의사는 소년에게 누구일까요?\n",
      "  Answer: 이 외과의사는 소년에게 누구일까요\n",
      "  Score: 0.0688\n",
      "\n",
      "Path 1:\n",
      "  Reasoning: 그 외과의사는 소년의 아버지입니다.\n",
      "  Answer: 그 외과의사는 소년의 아버지입니다\n",
      "  Score: 0.5734\n",
      "\n",
      "Path 2:\n",
      "  Reasoning: 이 외과의사는 소년의 아버지인 외과의사입니다.\n",
      "  Answer: 이 외과의사는 소년의 아버지인 외과의사입니다\n",
      "  Score: 0.6247\n",
      "\n",
      "Path 3:\n",
      "  Reasoning: 아이의 아버지인 외과의사가 말했어. \"난 수술 못해! 이 아이는 내 아들이라고!\" 이 외과의사는 소년에게 누구일까요?\n",
      "  Answer: 이 외과의사는 소년에게 누구일까요\n",
      "  Score: 0.0000\n",
      "\n",
      "Path 4:\n",
      "  Reasoning: 소년은 아버지인 외과의사의 아들입니다.\n",
      "  Answer: 소년은 아버지인 외과의사의 아들입니다\n",
      "  Score: 0.5413\n",
      "\n"
     ]
    }
   ],
   "source": [
    "prompt = '아이의 아버지인 외과의사가 말했어. \"난 수술 못해! 이 아이는 내 아들이라고!\" 이 외과의사는 소년에게 누구일까요?'\n",
    "# prompt = \"I have 3 apples, my dad has 2 more apples than me, how many apples do we have in total?\"\n",
    "result = decoder.search_cots(prompt)\n",
    "\n",
    "print(f\"Question: {result.question}\")\n",
    "for path in result.paths:\n",
    "    print(f\"Path {path.num_path}:\")\n",
    "    print(f\"  Reasoning: {path.reasoning_text}\")\n",
    "    print(f\"  Answer: {path.answer_span}\")\n",
    "    print(f\"  Score: {path.score:.4f}\")\n",
    "    print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "121"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "reasoning = \"I have 3 apples, my dad has 2 more apples than me, so my dad has 3 apples more than me. So, we have 3 apples + 3 apples = 6 apples.\"\n",
    "answer_span = ' 6 apples'\n",
    "reasoning.rfind(answer_span)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "vllm",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
